package util

import (
	"fmt"
	"github.com/go-redis/redis"
	"longmen/server/config/global"
	"longmen/server/internal/services/nami/utils"
	_const "longmen/server/pkg/common/helper/const"
	"strconv"
	"sync"
	"time"
)

type RequestLimitService struct {
	Interval time.Duration
	MaxCount int
	Lock     sync.Mutex
	ReqCount int
}

func NewRequestLimitService(interval time.Duration, maxCnt int) *RequestLimitService {
	reqLimit := &RequestLimitService{
		Interval: interval,
		MaxCount: maxCnt,
	}

	go func() {
		ticker := time.NewTicker(interval)
		for {
			<-ticker.C
			reqLimit.Lock.Lock()
			reqLimit.ReqCount = 0
			reqLimit.Lock.Unlock()
		}
	}()

	return reqLimit
}

func (reqLimit *RequestLimitService) Increase() {
	reqLimit.Lock.Lock()
	defer reqLimit.Lock.Unlock()

	reqLimit.ReqCount += 1
}

func (reqLimit *RequestLimitService) IsAvailable() bool {
	reqLimit.Lock.Lock()
	defer reqLimit.Lock.Unlock()

	return reqLimit.ReqCount < reqLimit.MaxCount
}

// 阻止短时间大量访问
func NeedIPRateLimit(ip string, maxReqLimit int) bool {
	windowSize, key := time.Second, fmt.Sprintf(_const.EsSeoIPLimit, ip)
	// 获取当前时间戳（豪秒级）
	now := time.Now()
	nowMilli := now.UnixMilli()
	// 滑动窗口开始结束时间戳
	windowSt := strconv.FormatInt(now.Add(-windowSize).UnixMilli(), 10)
	windowEnd := strconv.FormatInt(nowMilli, 10)

	// 获取指定时间窗口内的元素个数
	count, err := global.GetRds().ZCount(key, windowSt, windowEnd).Result()
	if err != nil {
		return false
	}
	needLimit := count > int64(maxReqLimit)

	// 不需要限流的访问, 记录当前访问记录
	if !needLimit {
		global.GetRds().ZAdd(key, redis.Z{
			Member: float64(nowMilli),
			Score:  float64(nowMilli),
		})
	}

	utils.SafeGo(func() {
		// 删除窗口外的元素
		global.GetRds().ZRemRangeByScore(key, "0", windowSt)
		// 设置过期时间一分钟
		global.GetRds().Expire(key, time.Second*60)
	})

	//log.Println(count, maxReqLimit)
	return needLimit
}
